import os
import random
import json
import webview
import threading
import hashlib
from model import train_model, recognize_word
from PIL import Image, ImageDraw
import io
import base64
import numpy as np
import re
import string
from wordlists import WORDLISTS, WORDLIST_NAMES

class WordGame:
    def __init__(self):
        self.window = None
        self.api = API()
        print("Training model...")
        self.qsvc_model, self.pca_model = train_model()
        print("Model training completed!")
        self.current_user = None
        self.users_dir = "user_data"
        os.makedirs(self.users_dir, exist_ok=True)
        self.users_file = os.path.join(self.users_dir, "users.json")
        if not os.path.exists(self.users_file):
            with open(self.users_file, 'w') as f:
                json.dump({}, f)

    def start(self):
        self.window = webview.create_window(
            title='单词默写大师',
            url=os.path.join(os.path.dirname(__file__), 'template', 'index.html'),
            js_api=self.api,
            width=1000,
            height=800,
            resizable=True,
        )
        webview.start(debug=False)


class API:
    def __init__(self):
        self.window = None
        self.draw_data = None
        self.current_draw_color = "black"
        self.history = []
        self.current_user = None
        self.users_dir = "user_data"
        self.users_file = os.path.join(self.users_dir, "users.json")

    def register_user(self, username, password):
        try:
            with open(self.users_file, 'r') as f:
                users = json.load(f)

            if username in users:
                return {"success": False, "message": "用户名已存在"}

            password_hash = hashlib.sha256(password.encode()).hexdigest()
            users[username] = {"password_hash": password_hash}
            with open(self.users_file, 'w') as f:
                json.dump(users, f, indent=4)

            user_data_file = os.path.join(self.users_dir, f"user_{username}.json")
            user_data = {
                "word_list": [],
                "current_word": None,
                "write_count": 0,
                "completed_words": [],
                "total_words": 0,
                "dark_mode": False,
                "canvas_image": None,
                "avatar": None,
                "birthday": None
            }
            with open(user_data_file, 'w') as f:
                json.dump(user_data, f, indent=4)

            return {"success": True, "message": "注册成功，请登录"}
        except Exception as e:
            return {"success": False, "message": f"注册失败: {str(e)}"}

    def login_user(self, username, password):
        try:
            with open(self.users_file, 'r') as f:
                users = json.load(f)

            if username not in users:
                return {"success": False, "message": "用户名不存在"}

            password_hash = hashlib.sha256(password.encode()).hexdigest()
            if users[username]["password_hash"] != password_hash:
                return {"success": False, "message": "密码错误"}

            self.current_user = username
            return {"success": True, "message": "登录成功"}
        except Exception as e:
            return {"success": False, "message": f"登录失败: {str(e)}"}

    def logout_user(self):
        try:
            if not self.current_user:
                return {"success": False, "message": "未登录"}

            self.current_user = None
            self.history = []
            return {"success": True, "message": "已退出登录"}
        except Exception as e:
            return {"success": False, "message": f"退出失败: {str(e)}"}

    def get_user_info(self):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        user_data = self.load_user_data()
        if not user_data:
            return {"success": False, "message": "用户数据加载失败"}

        return {
            "success": True,
            "username": self.current_user,
            "avatar": user_data.get("avatar"),
            "birthday": user_data.get("birthday"),
            "learned_words": len(set(tuple(word) for word in user_data.get("completed_words", [])))
        }

    def update_user_profile(self, avatar=None, birthday=None):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        try:
            user_data = self.load_user_data()
            if not user_data:
                return {"success": False, "message": "用户数据加载失败"}

            if avatar:
                user_data["avatar"] = avatar
            if birthday:
                user_data["birthday"] = birthday

            self.save_user_data(user_data)
            return {"success": True, "message": "个人信息更新成功"}
        except Exception as e:
            return {"success": False, "message": f"更新失败: {str(e)}"}

    def load_user_data(self):
        if not self.current_user:
            return None
        user_data_file = os.path.join(self.users_dir, f"user_{self.current_user}.json")
        if not os.path.exists(user_data_file):
            return None
        with open(user_data_file, 'r') as f:
            return json.load(f)

    def save_user_data(self, user_data):
        if not self.current_user:
            return
        user_data_file = os.path.join(self.users_dir, f"user_{self.current_user}.json")
        with open(user_data_file, 'w') as f:
            json.dump(user_data, f, indent=4)

    def load_custom_wordlist(self, file_content):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        try:
            words = []
            lines = file_content.split('\n')
            for line in lines:
                line = line.strip()
                if not line:
                    continue

                separators = r'[,:;\s]+'
                parts = re.split(separators, line)
                parts = [p.strip() for p in parts if p.strip()]

                if len(parts) >= 2:
                    chinese = parts[0]
                    english = ''.join(parts[1:])

                    english_lower = english.lower()
                    if english_lower and all(c in string.ascii_lowercase for c in english_lower):
                        words.append((chinese, english_lower))

            if not words:
                return {
                    "success": False,
                    "message": "未找到有效单词（必须包含且仅包含A-Z字母）"
                }

            user_data = self.load_user_data() or {}
            user_data["word_list"] = words
            user_data["total_words"] = len(words)
            user_data["current_word"] = random.choice(words)
            user_data["write_count"] = 1
            user_data["completed_words"] = []
            self.save_user_data(user_data)

            return {
                "success": True,
                "message": f"成功加载 {user_data['total_words']} 个单词",
                "currentWord": {
                    "chinese": user_data["current_word"][0],
                    "english": user_data["current_word"][1],
                    "progress": f"{min(user_data['write_count'], user_data['total_words'])}/{user_data['total_words']}"
                },
                "totalWords": user_data["total_words"]
            }
        except Exception as e:
            print(f"Error loading wordlist: {str(e)}")
            return {
                "success": False,
                "message": f"加载词书出错: {str(e)}"
            }

    def load_predefined_wordlist(self, wordlist_type):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        try:
            if wordlist_type not in WORDLISTS:
                return {"success": False, "message": "无效的词库类型"}

            words = WORDLISTS[wordlist_type]
            wordlist_name = WORDLIST_NAMES[wordlist_type]

            user_data = self.load_user_data() or {}
            user_data["word_list"] = words
            user_data["total_words"] = len(words)
            user_data["current_word"] = random.choice(words)
            user_data["write_count"] = 1
            user_data["completed_words"] = []
            self.save_user_data(user_data)

            return {
                "success": True,
                "message": f"成功加载 {wordlist_name}，共 {user_data['total_words']} 个单词",
                "wordlistName": wordlist_name,
                "currentWord": {
                    "chinese": user_data["current_word"][0],
                    "english": user_data["current_word"][1],
                    "progress": f"{min(user_data['write_count'], user_data['total_words'])}/{user_data['total_words']}"
                },
                "totalWords": user_data["total_words"]
            }
        except Exception as e:
            print(f"Error loading predefined wordlist: {str(e)}")
            return {
                "success": False,
                "message": f"加载词书出错: {str(e)}"
            }

    def get_next_word(self):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        user_data = self.load_user_data()
        if not user_data or not user_data["word_list"]:
            return {"success": False, "message": "请先加载词库"}

        completed_words = set(tuple(word) for word in user_data["completed_words"])
        if len(completed_words) == user_data["total_words"]:
            return {
                "completed": True,
                "message": "你已完成所有单词！"
            }

        remaining_words = [w for w in user_data["word_list"] if tuple(w) not in completed_words]
        user_data["current_word"] = random.choice(remaining_words)
        user_data["write_count"] = min(user_data["write_count"] + 1, user_data["total_words"])
        self.save_user_data(user_data)

        return {
            "completed": False,
            "currentWord": {
                "chinese": user_data["current_word"][0],
                "english": user_data["current_word"][1],
                "progress": f"{min(user_data['write_count'], user_data['total_words'])}/{user_data['total_words']}"
            }
        }

    def reset_wordlist(self):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        user_data = self.load_user_data()
        if not user_data or not user_data["word_list"]:
            return {"success": False, "message": "请先加载词库"}

        user_data["completed_words"] = []
        user_data["current_word"] = random.choice(user_data["word_list"])
        user_data["write_count"] = 1
        self.save_user_data(user_data)

        return {
            "success": True,
            "currentWord": {
                "chinese": user_data["current_word"][0],
                "english": user_data["current_word"][1],
                "progress": f"{min(user_data['write_count'], user_data['total_words'])}/{user_data['total_words']}"
            }
        }

    def toggle_dark_mode(self):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        user_data = self.load_user_data()
        user_data["dark_mode"] = not user_data.get("dark_mode", False)
        self.save_user_data(user_data)
        return {"darkMode": user_data["dark_mode"]}

    def save_drawing(self, image_data):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        try:
            image_data = re.sub('^data:image/.+;base64,', '', image_data)
            image_bytes = base64.b64decode(image_data)
            image = Image.open(io.BytesIO(image_bytes))

            user_data = self.load_user_data()
            user_data["canvas_image"] = base64.b64encode(image_bytes).decode('utf-8')
            self.save_user_data(user_data)

            return {"success": True}
        except Exception as e:
            print(f"Error saving drawing: {str(e)}")
            return {"success": False, "error": str(e)}

    def recognize_drawing(self):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        user_data = self.load_user_data()
        if not user_data or not user_data.get("canvas_image"):
            return {"success": False, "message": "No drawing to recognize"}

        try:
            image_bytes = base64.b64decode(user_data["canvas_image"])
            image = Image.open(io.BytesIO(image_bytes))
            if image.mode != 'L':
                image = image.convert('L')

            predicted_word = recognize_word(image, WordGame().qsvc_model, WordGame().pca_model)

            return {
                "success": True,
                "recognized": predicted_word
            }
        except Exception as e:
            print(f"Recognition error: {str(e)}")
            return {"success": False, "error": str(e)}

    def check_answer(self, user_answer):
        print(f"check_answer called with user_answer: {user_answer}")
        try:
            if not self.current_user:
                print("check_answer: current_user is None")
                return {"success": False, "message": "请先登录"}

            print(f"check_answer: current_user is {self.current_user}")
            user_data = self.load_user_data()
            if not user_data:
                print("check_answer: user_data is None")
                return {"success": False, "message": "用户数据加载失败，请重新登录"}
            print(f"check_answer: loaded user_data: {user_data}")

            if not user_data.get("current_word"):
                print("check_answer: current_word is None")
                return {"success": False, "message": "请先加载词库"}

            correct_answer = user_data["current_word"][1]
            is_correct = user_answer.strip().lower() == correct_answer.lower()
            print(f"check_answer: user_answer={user_answer}, correct_answer={correct_answer}, is_correct={is_correct}")

            if is_correct:
                completed_words = set(tuple(word) for word in user_data["completed_words"])
                completed_words.add(tuple(user_data["current_word"]))
                user_data["completed_words"] = list(completed_words)
                print("check_answer: saving user_data after correct answer")
                self.save_user_data(user_data)

            result = {
                "success": True,
                "correct": is_correct,
                "userAnswer": user_answer,
                "correctAnswer": correct_answer,
                "completedCount": len(user_data["completed_words"]),
                "totalCount": user_data["total_words"]
            }
            print(f"check_answer: returning result: {result}")
            return result
        except Exception as e:
            print(f"Error in check_answer: {str(e)}")
            return {"success": False, "message": f"答案检查失败: {str(e)}"}

    def get_current_status(self):
        if not self.current_user:
            return {
                "isLoggedIn": False,
                "hasWordList": False,
                "darkMode": False
            }

        user_data = self.load_user_data()
        if not user_data:
            return {
                "isLoggedIn": True,
                "hasWordList": False,
                "darkMode": False
            }

        return {
            "isLoggedIn": True,
            "hasWordList": bool(user_data.get("word_list")),
            "currentWord": user_data.get("current_word") and {
                "chinese": user_data["current_word"][0],
                "english": user_data["current_word"][1],
                "progress": f"{min(user_data['write_count'], user_data['total_words'])}/{user_data['total_words']}"
            },
            "darkMode": user_data.get("dark_mode", False)
        }

    def add_history(self, drawing_data):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        self.history.append(drawing_data)
        if len(self.history) > 20:
            self.history.pop(0)
        return {"historyLength": len(self.history)}

    def undo_drawing(self):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        if len(self.history) > 0:
            self.history.pop()
            if len(self.history) > 0:
                return {"success": True, "lastState": self.history[-1]}
            else:
                return {"success": True, "lastState": None}
        return {"success": False, "message": "No history to undo"}

    def set_draw_color(self, color):
        if not self.current_user:
            return {"success": False, "message": "请先登录"}

        self.current_draw_color = color
        return {"success": True, "color": color}


if __name__ == "__main__":
    game = WordGame()
    game.start()